package com.opages.mvc.adapter.consul.serviceregistry;

import com.ecwid.consul.v1.ConsulClient;
import com.ecwid.consul.v1.agent.model.NewService;
import com.opages.mvc.adapter.consul.model.ClusterConsulClient;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.concurrent.ConcurrentTaskScheduler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledFuture;

public class TtlScheduler {
    private final Map<String, ScheduledFuture> serviceHeartbeats = new ConcurrentHashMap<>();
    private final TaskScheduler scheduler = new ConcurrentTaskScheduler(Executors.newSingleThreadScheduledExecutor());
    private ClusterConsulClient clusterConsulClient;
    private long heartbeatInterval;

    public TtlScheduler(ClusterConsulClient clusterConsulClient,long heartbeatInterval) {
        this.clusterConsulClient = clusterConsulClient;
        this.heartbeatInterval = heartbeatInterval;
    }

    @Deprecated
    public void add(final NewService service) {
        add(service.getId());
    }

    /**
     * Add a service to the checks loop.
     * @param instanceId instance id
     */
    public void add(String instanceId) {
        ScheduledFuture task = this.scheduler.scheduleAtFixedRate(
                new ConsulHeartbeatTask(instanceId), heartbeatInterval);
        ScheduledFuture previousTask = this.serviceHeartbeats.put(instanceId, task);
        if (previousTask != null) {
            previousTask.cancel(true);
        }
    }

    public void remove(String instanceId) {
        ScheduledFuture task = this.serviceHeartbeats.get(instanceId);
        if (task != null) {
            task.cancel(true);
        }
        this.serviceHeartbeats.remove(instanceId);
    }

    private class ConsulHeartbeatTask implements Runnable {

        private String checkId;

        ConsulHeartbeatTask(String serviceId) {
            this.checkId = serviceId;
            if (!this.checkId.startsWith("service:")) {
                this.checkId = "service:" + this.checkId;
            }
        }

        @Override
        public void run() {
            for(ConsulClient client : clusterConsulClient.getClients()){
                client.agentCheckPass(this.checkId);
            }
        }

    }
}
